{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "import tensorflow as tf\n",
    "from PIL import Image\n",
    "from nets import nets_factory\n",
    "import numpy as np"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "# 字符集\n",
    "CHAR_SET = [str(i) for i in range(10)]\n",
    "CHAR_SET_LEN = len(CHAR_SET)\n",
    "# 训练集大小\n",
    "TRAIN_NUM = 5800\n",
    "# 批次大小\n",
    "BATCH_SIZE = 128\n",
    "# 迭代次数\n",
    "EPOCHES = 30\n",
    "# 循环次数\n",
    "LOOP_TIMES = EPOCHES*TRAIN_NUM//BATCH_SIZE\n",
    "# tf文件\n",
    "TFRECORD_FILE = 'captcha/train.tfrecord'\n",
    "\n",
    "# 初始学习率\n",
    "LEARNING_RATE = 0.001"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "# 读取数据的函数\n",
    "def read_and_decode(filename):\n",
    "    tf_queue = tf.train.string_input_producer([filename])\n",
    "    reader = tf.TFRecordReader()\n",
    "    _, serialized_example = reader.read(tf_queue)\n",
    "    features = tf.parse_single_example(serialized_example,\n",
    "                                       features={\n",
    "                                           'image':tf.FixedLenFeature([], tf.string),\n",
    "                                            'label0': tf.FixedLenFeature([], tf.int64),\n",
    "                                            'label1': tf.FixedLenFeature([], tf.int64),\n",
    "                                            'label2': tf.FixedLenFeature([], tf.int64),\n",
    "                                            'label3': tf.FixedLenFeature([], tf.int64),\n",
    "                                       })\n",
    "    image = tf.decode_raw(features['image'], tf.uint8)\n",
    "    image = tf.reshape(image,[224,224])\n",
    "    image = tf.cast(image, tf.float32) / 255.0\n",
    "    image = tf.subtract(image, 0.5)\n",
    "    image = tf.multiply(image, 2.0)\n",
    "\n",
    "    label0 = tf.cast(features['label0'], tf.int32)\n",
    "    label1 = tf.cast(features['label1'], tf.int32)\n",
    "    label2 = tf.cast(features['label2'], tf.int32)\n",
    "    label3 = tf.cast(features['label3'], tf.int32)\n",
    "    return image, label0, label1, label2, label3\n",
    "\n",
    "image, label0, label1, label2, label3 = read_and_decode(TFRECORD_FILE)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {
    "scrolled": false
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "WARNING:tensorflow:From <ipython-input-4-795731713057>:34: softmax_cross_entropy_with_logits (from tensorflow.python.ops.nn_ops) is deprecated and will be removed in a future version.\n",
      "Instructions for updating:\n",
      "\n",
      "Future major versions of TensorFlow will allow gradients to flow\n",
      "into the labels input on backprop by default.\n",
      "\n",
      "See @{tf.nn.softmax_cross_entropy_with_logits_v2}.\n",
      "\n",
      "Iter:0/1359 epoch:1,  Loss:10.469  Accuracy:0.15,0.15,0.09,0.12  Learning_rate:0.00100\n",
      "Iter:20/1359 epoch:1,  Loss:2.311  Accuracy:0.10,0.06,0.07,0.11  Learning_rate:0.00100\n",
      "Iter:40/1359 epoch:1,  Loss:2.302  Accuracy:0.12,0.09,0.12,0.14  Learning_rate:0.00100\n",
      "Iter:60/1359 epoch:2,  Loss:2.303  Accuracy:0.14,0.07,0.06,0.04  Learning_rate:0.00100\n",
      "Iter:80/1359 epoch:2,  Loss:2.299  Accuracy:0.14,0.10,0.17,0.09  Learning_rate:0.00100\n",
      "Iter:100/1359 epoch:3,  Loss:2.299  Accuracy:0.13,0.11,0.11,0.12  Learning_rate:0.00100\n",
      "Iter:120/1359 epoch:3,  Loss:2.303  Accuracy:0.16,0.11,0.08,0.09  Learning_rate:0.00100\n",
      "Iter:140/1359 epoch:4,  Loss:2.303  Accuracy:0.11,0.09,0.11,0.09  Learning_rate:0.00100\n",
      "Iter:160/1359 epoch:4,  Loss:2.301  Accuracy:0.16,0.15,0.12,0.12  Learning_rate:0.00100\n",
      "Iter:180/1359 epoch:4,  Loss:2.301  Accuracy:0.08,0.14,0.09,0.16  Learning_rate:0.00100\n",
      "Iter:200/1359 epoch:5,  Loss:2.302  Accuracy:0.09,0.12,0.16,0.12  Learning_rate:0.00100\n",
      "Iter:220/1359 epoch:5,  Loss:2.304  Accuracy:0.12,0.14,0.11,0.09  Learning_rate:0.00100\n",
      "Iter:240/1359 epoch:6,  Loss:2.303  Accuracy:0.11,0.18,0.07,0.12  Learning_rate:0.00100\n",
      "Iter:260/1359 epoch:6,  Loss:2.306  Accuracy:0.09,0.09,0.09,0.10  Learning_rate:0.00100\n",
      "Iter:280/1359 epoch:7,  Loss:2.307  Accuracy:0.08,0.11,0.09,0.11  Learning_rate:0.00100\n",
      "Iter:300/1359 epoch:7,  Loss:2.303  Accuracy:0.06,0.12,0.12,0.12  Learning_rate:0.00100\n",
      "Iter:320/1359 epoch:8,  Loss:2.305  Accuracy:0.11,0.05,0.16,0.09  Learning_rate:0.00050\n",
      "Iter:340/1359 epoch:8,  Loss:2.296  Accuracy:0.06,0.12,0.16,0.09  Learning_rate:0.00050\n",
      "Iter:360/1359 epoch:8,  Loss:2.248  Accuracy:0.14,0.10,0.16,0.16  Learning_rate:0.00050\n",
      "Iter:380/1359 epoch:9,  Loss:2.181  Accuracy:0.18,0.17,0.16,0.20  Learning_rate:0.00050\n",
      "Iter:400/1359 epoch:9,  Loss:2.080  Accuracy:0.14,0.28,0.21,0.13  Learning_rate:0.00050\n",
      "Iter:420/1359 epoch:10,  Loss:2.095  Accuracy:0.21,0.23,0.23,0.19  Learning_rate:0.00050\n",
      "Iter:440/1359 epoch:10,  Loss:1.957  Accuracy:0.23,0.23,0.22,0.23  Learning_rate:0.00050\n",
      "Iter:460/1359 epoch:11,  Loss:1.834  Accuracy:0.32,0.37,0.29,0.25  Learning_rate:0.00050\n",
      "Iter:480/1359 epoch:11,  Loss:1.814  Accuracy:0.38,0.25,0.17,0.28  Learning_rate:0.00050\n",
      "Iter:500/1359 epoch:12,  Loss:1.588  Accuracy:0.43,0.41,0.29,0.46  Learning_rate:0.00050\n",
      "Iter:520/1359 epoch:12,  Loss:1.453  Accuracy:0.49,0.39,0.41,0.46  Learning_rate:0.00050\n",
      "Iter:540/1359 epoch:12,  Loss:1.371  Accuracy:0.60,0.40,0.34,0.56  Learning_rate:0.00050\n",
      "Iter:560/1359 epoch:13,  Loss:1.378  Accuracy:0.55,0.41,0.44,0.51  Learning_rate:0.00050\n",
      "Iter:580/1359 epoch:13,  Loss:1.188  Accuracy:0.59,0.53,0.51,0.64  Learning_rate:0.00050\n",
      "Iter:600/1359 epoch:14,  Loss:1.046  Accuracy:0.68,0.59,0.50,0.63  Learning_rate:0.00050\n",
      "Iter:620/1359 epoch:14,  Loss:1.002  Accuracy:0.73,0.60,0.56,0.60  Learning_rate:0.00050\n",
      "Iter:640/1359 epoch:15,  Loss:0.863  Accuracy:0.84,0.58,0.59,0.75  Learning_rate:0.00050\n",
      "Iter:660/1359 epoch:15,  Loss:0.851  Accuracy:0.73,0.68,0.57,0.73  Learning_rate:0.00050\n",
      "Iter:680/1359 epoch:16,  Loss:0.810  Accuracy:0.79,0.62,0.60,0.75  Learning_rate:0.00025\n",
      "Iter:700/1359 epoch:16,  Loss:0.716  Accuracy:0.82,0.70,0.67,0.75  Learning_rate:0.00025\n",
      "Iter:720/1359 epoch:16,  Loss:0.565  Accuracy:0.89,0.77,0.71,0.80  Learning_rate:0.00025\n",
      "Iter:740/1359 epoch:17,  Loss:0.575  Accuracy:0.87,0.71,0.73,0.83  Learning_rate:0.00025\n",
      "Iter:760/1359 epoch:17,  Loss:0.591  Accuracy:0.81,0.77,0.73,0.81  Learning_rate:0.00025\n",
      "Iter:780/1359 epoch:18,  Loss:0.501  Accuracy:0.84,0.80,0.78,0.84  Learning_rate:0.00025\n",
      "Iter:800/1359 epoch:18,  Loss:0.549  Accuracy:0.84,0.75,0.77,0.85  Learning_rate:0.00025\n",
      "Iter:820/1359 epoch:19,  Loss:0.439  Accuracy:0.88,0.80,0.79,0.85  Learning_rate:0.00025\n",
      "Iter:840/1359 epoch:19,  Loss:0.364  Accuracy:0.84,0.77,0.86,0.93  Learning_rate:0.00025\n",
      "Iter:860/1359 epoch:19,  Loss:0.419  Accuracy:0.85,0.84,0.76,0.87  Learning_rate:0.00025\n",
      "Iter:880/1359 epoch:20,  Loss:0.423  Accuracy:0.83,0.88,0.78,0.90  Learning_rate:0.00025\n",
      "Iter:900/1359 epoch:20,  Loss:0.431  Accuracy:0.87,0.86,0.79,0.83  Learning_rate:0.00025\n",
      "Iter:920/1359 epoch:21,  Loss:0.427  Accuracy:0.91,0.82,0.79,0.86  Learning_rate:0.00025\n",
      "Iter:940/1359 epoch:21,  Loss:0.343  Accuracy:0.89,0.86,0.91,0.88  Learning_rate:0.00025\n",
      "Iter:960/1359 epoch:22,  Loss:0.354  Accuracy:0.92,0.85,0.85,0.89  Learning_rate:0.00025\n",
      "Iter:980/1359 epoch:22,  Loss:0.316  Accuracy:0.93,0.86,0.88,0.88  Learning_rate:0.00025\n",
      "Iter:1000/1359 epoch:23,  Loss:0.323  Accuracy:0.90,0.86,0.88,0.88  Learning_rate:0.00025\n",
      "Iter:1020/1359 epoch:23,  Loss:0.289  Accuracy:0.90,0.91,0.90,0.88  Learning_rate:0.00025\n",
      "Iter:1040/1359 epoch:23,  Loss:0.328  Accuracy:0.96,0.86,0.86,0.87  Learning_rate:0.00025\n",
      "Iter:1060/1359 epoch:24,  Loss:0.264  Accuracy:0.94,0.86,0.91,0.88  Learning_rate:0.00013\n",
      "Iter:1080/1359 epoch:24,  Loss:0.268  Accuracy:0.94,0.88,0.91,0.90  Learning_rate:0.00013\n",
      "Iter:1100/1359 epoch:25,  Loss:0.267  Accuracy:0.92,0.85,0.89,0.90  Learning_rate:0.00013\n",
      "Iter:1120/1359 epoch:25,  Loss:0.278  Accuracy:0.95,0.85,0.87,0.89  Learning_rate:0.00013\n",
      "Iter:1140/1359 epoch:26,  Loss:0.229  Accuracy:0.94,0.91,0.91,0.94  Learning_rate:0.00013\n",
      "Iter:1160/1359 epoch:26,  Loss:0.181  Accuracy:0.97,0.92,0.92,0.93  Learning_rate:0.00013\n",
      "Iter:1180/1359 epoch:27,  Loss:0.204  Accuracy:0.94,0.88,0.94,0.94  Learning_rate:0.00013\n",
      "Iter:1200/1359 epoch:27,  Loss:0.211  Accuracy:0.94,0.94,0.90,0.93  Learning_rate:0.00013\n",
      "Iter:1220/1359 epoch:27,  Loss:0.238  Accuracy:0.97,0.89,0.89,0.92  Learning_rate:0.00013\n",
      "Iter:1240/1359 epoch:28,  Loss:0.183  Accuracy:0.92,0.92,0.91,0.98  Learning_rate:0.00013\n",
      "Iter:1260/1359 epoch:28,  Loss:0.220  Accuracy:0.95,0.91,0.88,0.92  Learning_rate:0.00013\n",
      "Iter:1280/1359 epoch:29,  Loss:0.225  Accuracy:0.93,0.92,0.91,0.94  Learning_rate:0.00013\n",
      "Iter:1300/1359 epoch:29,  Loss:0.157  Accuracy:0.98,0.95,0.92,0.95  Learning_rate:0.00013\n",
      "Iter:1320/1359 epoch:30,  Loss:0.173  Accuracy:0.95,0.91,0.93,0.91  Learning_rate:0.00013\n",
      "Iter:1340/1359 epoch:30,  Loss:0.155  Accuracy:0.93,0.95,0.94,0.98  Learning_rate:0.00013\n"
     ]
    }
   ],
   "source": [
    "image_batch, label0_batch, label1_batch, label2_batch, label3_batch = tf.train.shuffle_batch(\n",
    "    [image, label0, label1, label2, label3],\n",
    "    batch_size=BATCH_SIZE,\n",
    "    capacity=10000,\n",
    "    min_after_dequeue=2000,\n",
    "    num_threads=1\n",
    ")\n",
    "\n",
    "\n",
    "# 定义网络结构\n",
    "train_network_fn = nets_factory.get_network_fn(\n",
    "    'alexnet_v2_captcha_multi',\n",
    "    num_classes=CHAR_SET_LEN,\n",
    "    weight_decay=0.0005,\n",
    "    is_training=True\n",
    ")\n",
    "\n",
    "# 网络\n",
    "x = tf.placeholder(tf.float32,[None,224,224])\n",
    "y0 = tf.placeholder(tf.float32,[None])\n",
    "y1 = tf.placeholder(tf.float32,[None])\n",
    "y2 = tf.placeholder(tf.float32,[None])\n",
    "y3 = tf.placeholder(tf.float32,[None])\n",
    "lr = tf.Variable(LEARNING_RATE, dtype=tf.float32)\n",
    "\n",
    "X = tf.reshape(x,[BATCH_SIZE,224,224,1])\n",
    "logits0, logits1, logits2, logits3, end_points = train_network_fn(X)\n",
    "\n",
    "one_hot_label0 = tf.one_hot(indices=tf.cast(y0,tf.int32), depth=CHAR_SET_LEN)\n",
    "one_hot_label1 = tf.one_hot(indices=tf.cast(y1,tf.int32), depth=CHAR_SET_LEN)\n",
    "one_hot_label2 = tf.one_hot(indices=tf.cast(y2,tf.int32), depth=CHAR_SET_LEN)\n",
    "one_hot_label3 = tf.one_hot(indices=tf.cast(y3,tf.int32), depth=CHAR_SET_LEN)\n",
    "\n",
    "loss0 = tf.reduce_mean(tf.nn.softmax_cross_entropy_with_logits(labels=one_hot_label0,logits=logits0))\n",
    "loss1 = tf.reduce_mean(tf.nn.softmax_cross_entropy_with_logits(labels=one_hot_label1,logits=logits1))\n",
    "loss2 = tf.reduce_mean(tf.nn.softmax_cross_entropy_with_logits(labels=one_hot_label2,logits=logits2))\n",
    "loss3 = tf.reduce_mean(tf.nn.softmax_cross_entropy_with_logits(labels=one_hot_label3,logits=logits3))\n",
    "\n",
    "total_loss = (loss0 + loss1 + loss2 + loss3)/4.0\n",
    "optimizer = tf.train.AdamOptimizer(learning_rate=lr).minimize(total_loss)\n",
    "\n",
    "# 计算准确率\n",
    "correct_pre0 = tf.equal(tf.argmax(one_hot_label0, 1), tf.argmax(logits0, 1))\n",
    "accuracy0 = tf.reduce_mean(tf.cast(correct_pre0, tf.float32))\n",
    "\n",
    "correct_pre1 = tf.equal(tf.argmax(one_hot_label1, 1), tf.argmax(logits1, 1))\n",
    "accuracy1 = tf.reduce_mean(tf.cast(correct_pre1, tf.float32))\n",
    "\n",
    "correct_pre2 = tf.equal(tf.argmax(one_hot_label2, 1), tf.argmax(logits2, 1))\n",
    "accuracy2 = tf.reduce_mean(tf.cast(correct_pre2, tf.float32))\n",
    "\n",
    "correct_pre3 = tf.equal(tf.argmax(one_hot_label3, 1), tf.argmax(logits3, 1))\n",
    "accuracy3 = tf.reduce_mean(tf.cast(correct_pre3, tf.float32))\n",
    "\n",
    "saver = tf.train.Saver()\n",
    "\n",
    "\n",
    "with tf.Session() as sess:\n",
    "    sess.run(tf.global_variables_initializer())\n",
    "    coord = tf.train.Coordinator()\n",
    "    threads = tf.train.start_queue_runners(sess=sess,coord=coord)\n",
    "\n",
    "    i_epoch = 0\n",
    "\n",
    "    for i in range(LOOP_TIMES):\n",
    "        b_image, b_label0, b_label1, b_label2, b_label3 = sess.run(\n",
    "            [image_batch, label0_batch, label1_batch, label2_batch, label3_batch])\n",
    "        sess.run(optimizer, feed_dict={x:b_image,\n",
    "                                       y0:b_label0,\n",
    "                                       y1:b_label1,\n",
    "                                       y2:b_label2,\n",
    "                                       y3:b_label3})\n",
    "        i_epoch_new = i//(5800/BATCH_SIZE) + 1\n",
    "        if i_epoch != i_epoch_new:\n",
    "            i_epoch = i_epoch_new\n",
    "            if i_epoch%8 == 0:\n",
    "                sess.run(tf.assign(lr,lr*0.5))\n",
    "\n",
    "        if i%20 == 0:\n",
    "            acc0, acc1, acc2, acc3, loss_ = sess.run([accuracy0, accuracy1,accuracy2, accuracy3,total_loss],\n",
    "                                                     feed_dict={x:b_image,\n",
    "                                                               y0:b_label0,\n",
    "                                                               y1:b_label1,\n",
    "                                                               y2:b_label2,\n",
    "                                                               y3:b_label3})\n",
    "            learning_rate = sess.run(lr)\n",
    "            print(\"Iter:%d/%d epoch:%d,  Loss:%.3f  Accuracy:%.2f,%.2f,%.2f,%.2f  Learning_rate:%.5f\" % (\n",
    "                i, LOOP_TIMES, i_epoch, loss_, acc0, acc1, acc2, acc3, learning_rate))\n",
    "        if acc0>0.9 and acc1>0.9 and acc2>0.9 and acc3>0.9 and i==LOOP_TIMES-1:\n",
    "            saver.save(sess,'captcha/model/crack_captcha.model',global_step=i)\n",
    "\n",
    "    coord.request_stop()\n",
    "    coord.join(threads)"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.6.6"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}
